//routines for finding
// * MNT curves with embedding degree 6
// * Freeman curves (which have embedding degree 10)

#include <stdio.h>
#include <stdlib.h>
#include <gmp.h>
#include "pbc_darray.h"
#include "pbc_mnt.h"
#include "pbc_memory.h"
#include "pbc_utils.h"

struct pell_solution_s {
    int count;
    mpz_t minx; //minimal solution of x^2 - Dy^2 = 1
    mpz_t miny;
    mpz_t *x;
    mpz_t *y;
};
typedef struct pell_solution_s pell_solution_t[1];
typedef struct pell_solution_s *pell_solution_ptr;

static void general_pell(pell_solution_t ps, mpz_t D, int N)
//solves x^2 - Dy^2 = N
//D not a square
//(for square D, observe (x+Dy)(x-Dy) = N and look at factors of N)
//TODO: brute force for small D
{
    darray_t L;
    int i, f, n, sgnN = N > 0 ? 1 : -1;

    //find square factors of N
    darray_t listf;
    darray_init(listf);

    f = 1;
    for (;;) {
	n = f * f;
	if (n > abs(N)) break;
	if (!(abs(N) % n)) {
	    darray_append(listf, int_to_voidp(f));
	}
	f++;
    }

    //a0, twice_a0 don't change once initialized
    //a1 is a_i every iteration
    //P0, P1 become P_{i-1}, P_i every iteration
    //similarly for Q0, Q1
    mpz_t a0, twice_a0, a1;
    mpz_t P0, P1;
    mpz_t Q0, Q1;
    //variables to compute the convergents
    mpz_t p0, p1, pnext;
    mpz_t q0, q1, qnext;

    int d;

    darray_t listp, listq;
    mpz_ptr zptr;

    mpz_init(a0);
    mpz_init(twice_a0);
    mpz_init(a1);
    mpz_init(P0); mpz_init(P1);
    mpz_init(Q0); mpz_init(Q1);
    mpz_init(p0); mpz_init(p1); mpz_init(pnext);
    mpz_init(q0); mpz_init(q1); mpz_init(qnext);

    darray_init(L);

    darray_init(listp);
    darray_init(listq);

    mpz_sqrt(a0, D);
    mpz_set_ui(P0, 0);
    mpz_set_ui(Q0, 1);

    mpz_set(P1, a0);
    mpz_mul(Q1, a0, a0);
    mpz_sub(Q1, D, Q1);
    mpz_add(a1, a0, P1);
    mpz_tdiv_q(a1, a1, Q1);

    mpz_add(twice_a0, a0, a0);

    mpz_set(p0, a0);
    mpz_set_ui(q0, 1);
    mpz_mul(p1, a0, a1);
    mpz_add_ui(p1, p1, 1);
    mpz_set(q1, a1);

    d = -1;
    for(;;) {
	if (d == sgnN) {
	    for (i=0; i<listf->count; i++) {
		f = (int) listf->item[i];
		if (!mpz_cmp_ui(Q1, abs(N) / (f * f))) {
//element_printf("found %Zd, %Zd, %d\n", p0, q0, f);
		    zptr = (mpz_ptr) pbc_malloc(sizeof(mpz_t));
		    mpz_init(zptr);
		    mpz_set(zptr, p0);
		    mpz_mul_ui(zptr, p0, f);
		    darray_append(listp, zptr);
		    zptr = (mpz_ptr) pbc_malloc(sizeof(mpz_t));
		    mpz_init(zptr);
		    mpz_set(zptr, q0);
		    mpz_mul_ui(zptr, q0, f);
		    darray_append(listq, zptr);
		}
	    }
	}

	if (!mpz_cmp(twice_a0, a1) && d == 1) break;
	//compute more of the continued fraction expansion
	mpz_set(P0, P1);
	mpz_mul(P1, a1, Q1);
	mpz_sub(P1, P1, P0);
	mpz_set(Q0, Q1);
	mpz_mul(Q1, P1, P1);
	mpz_sub(Q1, D, Q1);
	mpz_divexact(Q1, Q1, Q0);
	mpz_add(a1, a0, P1);
	mpz_tdiv_q(a1, a1, Q1);

	//compute next convergent
	mpz_mul(pnext, a1, p1);
	mpz_add(pnext, pnext, p0);
	mpz_set(p0, p1);
	mpz_set(p1, pnext);

	mpz_mul(qnext, a1, q1);
	mpz_add(qnext, qnext, q0);
	mpz_set(q0, q1);
	mpz_set(q1, qnext);
	d = -d;
    }
    darray_clear(listf);

    mpz_init(ps->minx);
    mpz_init(ps->miny);
    mpz_set(ps->minx, p0);
    mpz_set(ps->miny, q0);
    n = listp->count;
    ps->count = n;
    if (n) {
	ps->x = (mpz_t *) pbc_malloc(sizeof(mpz_t) * n);
	ps->y = (mpz_t *) pbc_malloc(sizeof(mpz_t) * n);
	for (i=0; i<n; i++) {
	    mpz_init(ps->x[i]);
	    mpz_init(ps->y[i]);
	    mpz_set(ps->x[i], (mpz_ptr) listp->item[i]);
	    mpz_set(ps->y[i], (mpz_ptr) listq->item[i]);
	}
    }

    mpz_clear(a0);
    mpz_clear(twice_a0);
    mpz_clear(a1);
    mpz_clear(P0); mpz_clear(P1);
    mpz_clear(Q0); mpz_clear(Q1);
    mpz_clear(p0); mpz_clear(p1); mpz_clear(pnext);
    mpz_clear(q0); mpz_clear(q1); mpz_clear(qnext);
}

static void pell_solution_clear(pell_solution_t ps)
{
    int i, n = ps->count;

    if (n) {
	for (i=0; i<n; i++) {
	    mpz_clear(ps->x[i]);
	    mpz_clear(ps->y[i]);
	}
	pbc_free(ps->x);
	pbc_free(ps->y);
    }
    mpz_clear(ps->minx);
    mpz_clear(ps->miny);
}

void cm_info_init(cm_info_t cm)
{
    mpz_init(cm->q);
    mpz_init(cm->r);
    mpz_init(cm->h);
    mpz_init(cm->n);
}

void cm_info_clear(cm_info_t cm)
{
    mpz_clear(cm->q);
    mpz_clear(cm->r);
    mpz_clear(cm->h);
    mpz_clear(cm->n);
}

static int mnt_step2(darray_ptr L, unsigned int D, mpz_t U)
{
    int d;
    mpz_t n, l, q;
    mpz_t p;
    mpz_t r, cofac;
    cm_info_ptr cm;

    mpz_init(l);
    mpz_mod_ui(l, U, 6);
    if (!mpz_cmp_ui(l, 1)) {
	mpz_sub_ui(l, U, 1);
	d = 1;
    } else if (!mpz_cmp_ui(l, 5)) {
	mpz_add_ui(l, U, 1);
	d = -1;
    } else {
	mpz_clear(l);
	return 1;
    }

    mpz_divexact_ui(l, l, 3);
    mpz_init(q);

    mpz_mul(q, l, l);
    mpz_add_ui(q, q, 1);
    if (!mpz_probab_prime_p(q, 10)) {
	mpz_clear(q);
	mpz_clear(l);
	return 1;
    }

    mpz_init(n);
    if (d < 0) {
	mpz_sub(n, q, l);
    } else {
	mpz_add(n, q, l);
    }


    mpz_init(p);
    mpz_init(r);
    mpz_init(cofac);
{
    mpz_set_ui(cofac, 1);
    mpz_set(r, n);
    mpz_set_ui(p, 2);
    if (!mpz_probab_prime_p(r, 10)) for(;;) {
	if (mpz_divisible_p(r, p)) do {
	    mpz_mul(cofac, cofac, p);
	    mpz_divexact(r, r, p);
	} while (mpz_divisible_p(r, p));
	if (mpz_probab_prime_p(r, 10)) break;
	//TODO: use a table of primes instead?
	mpz_nextprime(p, p);
	if (mpz_sizeinbase(p, 2) > 16) {
	    //printf("has 16+ bit factor\n");
	    mpz_clear(r);
	    mpz_clear(p);
	    mpz_clear(cofac);
	    mpz_clear(q);
	    mpz_clear(l);
	    mpz_clear(n);
	    return 1;
	}
    }
}

    cm = pbc_malloc(sizeof(cm_info_t));
    cm_info_init(cm);
    cm->k = 6;
    cm->D = D;
    mpz_set(cm->q, q);
    mpz_set(cm->r, r);
    mpz_set(cm->h, cofac);
    mpz_set(cm->n, n);
    darray_append(L, cm);

    mpz_clear(cofac);
    mpz_clear(r);
    mpz_clear(p);
    mpz_clear(q);
    mpz_clear(l);
    mpz_clear(n);
    return 0;
}

int find_mnt6_curve(darray_t L, unsigned int D, unsigned int bitlimit)
{
    mpz_t D3;
    mpz_t t0, t1, t2;

    int found_count = 0;

    mpz_init(D3);
    mpz_set_ui(D3, D * 3);

    if (mpz_perfect_square_p(D3)) {
    //(the only squares that differ by 8 are 1 and 9,
    //which we get if U=V=1, D=3, but then l is not an integer)
	mpz_clear(D3);
	return 0;
    }

    mpz_init(t0);
    mpz_init(t1);
    mpz_init(t2);

    pell_solution_t ps;

    general_pell(ps, D3, -8);

    int i, n;
    n = ps->count;
    if (n) for (;;) {
	for (i=0; i<n; i++) {
	    //element_printf("%Zd, %Zd\n", ps->x[i], ps->y[i]);
	    if (!mnt_step2(L, D, ps->x[i])) found_count++;
	    //compute next solution as follows
	    //if p, q is current solution
	    //compute new solution p', q' via
	    //(p + q sqrt{3D})(t + u sqrt{3D}) = p' + q' sqrt(3D)
	    //where t, u is min. solution to Pell equation
	    mpz_mul(t0, ps->minx, ps->x[i]);
	    mpz_mul(t1, ps->miny, ps->y[i]);
	    mpz_mul(t1, t1, D3);
	    mpz_add(t0, t0, t1);
	    if (2 * mpz_sizeinbase(t0, 2) > bitlimit + 10) goto toobig;
	    mpz_mul(t2, ps->minx, ps->y[i]);
	    mpz_mul(t1, ps->miny, ps->x[i]);
	    mpz_add(t2, t2, t1);
	    mpz_set(ps->x[i], t0);
	    mpz_set(ps->y[i], t2);
	}
    }
toobig:

    pell_solution_clear(ps);
    mpz_clear(t0);
    mpz_clear(t1);
    mpz_clear(t2);
    mpz_clear(D3);
    return found_count;
}

static int freeman_step2(darray_ptr L, unsigned int D, mpz_t U)
{
    mpz_t n, x, q;
    mpz_t p;
    mpz_t r, cofac;
    cm_info_ptr cm;

    mpz_init(x);
    mpz_mod_ui(x, U, 15);
    if (!mpz_cmp_ui(x, 5)) {
	mpz_sub_ui(x, U, 5);
    } else if (!mpz_cmp_ui(x, 10)) {
	mpz_add_ui(x, U, 5);
    } else {
	//TODO: this code should never be reached
	mpz_clear(x);
	return 1;
    }

    mpz_divexact_ui(x, x, 15);
    mpz_init(q);
    mpz_init(r);

    //q = 25x^4 + 25x^3 + 25x^2 + 10x + 3
    mpz_mul(r, x, x);
    mpz_add(q, x, x);
    mpz_mul_ui(r, r, 5);
    mpz_add(q, q, r);
    mpz_mul(r, r, x);
    mpz_add(q, q, r);
    mpz_mul(r, r, x);
    mpz_add(q, q, r);
    mpz_mul_ui(q, q, 5);
    mpz_add_ui(q, q, 3);

    if (!mpz_probab_prime_p(q, 10)) {
	mpz_clear(q);
	mpz_clear(r);
	mpz_clear(x);
	return 1;
    }

    //t = 10x^2 + 5x + 3
    //n = q - t + 1
    mpz_init(n);

    mpz_mul_ui(n, x, 5);
    mpz_mul(r, n, x);
    mpz_add(r, r, r);
    mpz_add(n, n, r);
    mpz_sub(n, q, n);
    mpz_sub_ui(n, n, 2);

    mpz_init(p);
    mpz_init(cofac);
{
    mpz_set_ui(cofac, 1);
    mpz_set(r, n);
    mpz_set_ui(p, 2);
    if (!mpz_probab_prime_p(r, 10)) for(;;) {
	if (mpz_divisible_p(r, p)) do {
	    mpz_mul(cofac, cofac, p);
	    mpz_divexact(r, r, p);
	} while (mpz_divisible_p(r, p));
	if (mpz_probab_prime_p(r, 10)) break;
	//TODO: use a table of primes instead?
	mpz_nextprime(p, p);
	if (mpz_sizeinbase(p, 2) > 16) {
	    //printf("has 16+ bit factor\n");
	    mpz_clear(r);
	    mpz_clear(p);
	    mpz_clear(cofac);
	    mpz_clear(q);
	    mpz_clear(x);
	    mpz_clear(n);
	    return 1;
	}
    }
}

    cm = pbc_malloc(sizeof(cm_info_t));
    cm_info_init(cm);
    cm->k = 10;
    cm->D = D;
    mpz_set(cm->q, q);
    mpz_set(cm->r, r);
    mpz_set(cm->h, cofac);
    mpz_set(cm->n, n);
    darray_append(L, cm);

    mpz_clear(cofac);
    mpz_clear(r);
    mpz_clear(p);
    mpz_clear(q);
    mpz_clear(x);
    mpz_clear(n);
    return 0;
}

int find_freeman_curve(darray_t L, unsigned int D, unsigned int bitlimit)
{
    mpz_t D15;
    mpz_t t0, t1, t2;

    int found_count = 0;

    mpz_init(D15);
    //mpz_set_ui(D15, D * 15);
    mpz_set_ui(D15, D);
    mpz_mul_ui(D15, D15, 15);

    mpz_init(t0);
    mpz_init(t1);
    mpz_init(t2);

    if (mpz_perfect_square_p(D15)) {
	return 0;
    }

    pell_solution_t ps;

    general_pell(ps, D15, -20);

    int i, n;
    n = ps->count;
    if (n) for (;;) {
	for (i=0; i<n; i++) {
	    //gmp_printf("%Zd, %Zd\n", ps->x[i], ps->y[i]);
	    if (!freeman_step2(L, D, ps->x[i])) found_count++;
	    //compute next solution as follows
	    //if p, q is current solution
	    //compute new solution p', q' via
	    //(p + q sqrt{15D})(t + u sqrt{15D}) = p' + q' sqrt(15D)
	    //where t, u is min. solution to Pell equation
	    mpz_mul(t0, ps->minx, ps->x[i]);
	    mpz_mul(t1, ps->miny, ps->y[i]);
	    mpz_mul(t1, t1, D15);
	    mpz_add(t0, t0, t1);
	    if (2 * mpz_sizeinbase(t0, 2) > bitlimit + 10) goto toobig;
	    mpz_mul(t2, ps->minx, ps->y[i]);
	    mpz_mul(t1, ps->miny, ps->x[i]);
	    mpz_add(t2, t2, t1);
	    mpz_set(ps->x[i], t0);
	    mpz_set(ps->y[i], t2);
	}
    }
toobig:

    pell_solution_clear(ps);
    mpz_clear(t0);
    mpz_clear(t1);
    mpz_clear(t2);
    mpz_clear(D15);
    return found_count;
}
